package ai.neuralnet.builders;

import ai.neuralnet.NNBuilder;
import ai.neuralnet.NeuralNetwork;
import ai.neuralnet.TNeuron;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;

/**
 * Builds a fully connected multilayer network, one can specify number of input nodes, output nodes, hidden layers and nodes per hidden layer
 *
 * @param <N>  type of Neuron object to build
 * @param <NN> type of NeuralNetwork object to build
 */
public class ConnectedMultilayerBuilder<N extends TNeuron, NN extends NeuralNetwork<N>> extends NNBuilder<N, NN>
{

    private int hiddenLayers;

    private int hiddenLayerSize;


    public ConnectedMultilayerBuilder(Class<N> neuronClass, Class<NN> networkClass)
    {
        super(neuronClass, networkClass);
    }

    @Override
    public NN build()
    {
        NN network = makeNeuralNetwork();

        //create input layer neurons
        N[] inputNeurons = makeNeuronArray(inputCount);
        for (int i = 0; i < inputCount; i++)
        {

            inputNeurons[i] = makeNeuron(1);
            inputNeurons[i].setIsInputNeuron(true);

        }
        //add them to the network
        network.setInputNeurons(inputNeurons);

        //create hidden layers
        List<N> prevLayer = Arrays.asList(inputNeurons);
        for (int i = 0; i < hiddenLayers; i++)
        {
            //create new hidden layer
            ArrayList<N> layer = new ArrayList<N>();
            for (int j = 0; j < hiddenLayerSize; j++)
            {
                //set inputs of new neuron to the input layer neurons
                N n = makeNeuron(prevLayer.size());
                for (N iNeuron : prevLayer)
                {
                    n.addInputNeuron(iNeuron);
                    iNeuron.addOutputNeuron(n);
                }
                layer.add(n);

            }

            //add layer to network
            network.addHiddenNeuron(layer);

            prevLayer = layer;
        }

        //create output layer neurons
        ArrayList<N> outputNeurons = new ArrayList<N>();
        for (int i = 0; i < outputCount; i++)
        {
            N n = makeNeuron(prevLayer.size());
            //set inputs of new GDNeuron to last hidden layer
            for (N iNeuron : prevLayer)
            {
                n.addInputNeuron(iNeuron);
                iNeuron.addOutputNeuron(n);
            }
            outputNeurons.add(n);
        }

        N[] neurons = makeNeuronArray(outputNeurons.size());
        outputNeurons.toArray(neurons);
        network.setOutputNeurons(neurons);

        return network;
    }


    public void setHiddenLayers(int hiddenLayers)
    {
        this.hiddenLayers = hiddenLayers;
    }

    public void setHiddenLayerSize(int hiddenLayerSize)
    {
        this.hiddenLayerSize = hiddenLayerSize;
    }
}
